
from core.layer import Layer
from core.cuda import cuda_module
from implement.variables.parameter import Parameter
from utils.functions_collect import batch_norm


class BatchNorm(Layer):
    def __init__(self):
        super().__init__()
        # `.avg_mean` and `.avg_var` are `Parameter` objects, so they will be
        # saved to a file (using `save_weights()`).
        # But they don't need grads, so they're just used as `ndarray`.
        self.avg_mean = Parameter(None, name='avg_mean')
        self.avg_var = Parameter(None, name='avg_var')
        self.gamma = Parameter(None, name='gamma')
        self.beta = Parameter(None, name='beta')

    def _init_params(self, x):
        xp = cuda_module
        D = x.shape[1]
        # 确保 avg_mean, avg_var, gamma, beta 都有数据
        self.avg_mean.data = xp.zeros(D, dtype=x.dtype) if self.avg_mean.data is None else self.avg_mean.data
        self.avg_var.data = xp.ones(D, dtype=x.dtype) if self.avg_var.data is None else self.avg_var.data
        self.gamma.data = xp.ones(D, dtype=x.dtype) if self.gamma.data is None else self.gamma.data
        self.beta.data = xp.zeros(D, dtype=x.dtype) if self.beta.data is None else self.beta.data

    def __call__(self, x):
        """
        对输入数据进行 Batch Normalization 处理。

        Parameters:
            x: 输入数据。

        Returns:
            y: Batch Normalization 处理后的数据。
        """
        # 如果平均均值为空，初始化参数
        if self.avg_mean.data is None:
            self._init_params(x)
        # 调用 batch_norm 函数进行处理
        return batch_norm(x, self.gamma, self.beta, self.avg_mean.data,
                          self.avg_var.data)
